#!/usr/bin/env python3
#
# Copyright 2019-2021 Axel Huebl, Luca Fedeli, Maxence Thevenet
#
#
# This file is part of WarpX.
#
# License: BSD-3-Clause-LBNL

"""
This python script reads ionization tables in atomic_data.txt (generated from
the NIST website) and extracts ionization levels into C++ file
IonizationEnergiesTable.H, which contains tables + metadata.
"""

import os
import re

import numpy as np

filename = os.path.join(".", "atomic_data.txt")
with open(filename) as f:
    text_data = f.read()

# Read full table from file and get names, atomic numbers and offsets
# position in table of ionization energies for all species
regex_command = (
    "\n\s+(\d+)\s+\|\s+([A-Z]+[a-z]*)\s+\w+\s+\|\s+\+*(\d+)\s+\|\s+\(*\[*(\d+\.*\d*)"
)
list_of_tuples = re.findall(regex_command, text_data)
ion_atom_numbers = [int(i) for i in list(dict.fromkeys([x[0] for x in list_of_tuples]))]
ion_names = list(dict.fromkeys([x[1] for x in list_of_tuples]))
ion_offsets = np.concatenate(([0], np.cumsum(np.array(ion_atom_numbers)[:-1])), axis=0)

# Head of CPP file
cpp_string = "// This script was automatically generated!\n"
cpp_string += "// Edit dev/Source/Utils/Physics/write_atomic_data_cpp.py instead!\n\n"
cpp_string += "#ifndef WARPX_UTILS_PHYSICS_IONIZATION_TABLE_H_\n"
cpp_string += "#define WARPX_UTILS_PHYSICS_IONIZATION_TABLE_H_\n\n"
cpp_string += "#include <AMReX_REAL.H>\n\n"
cpp_string += "#include <map>\n"
cpp_string += "#include <string>\n\n"
cpp_string += "namespace utils::physics\n"
cpp_string += "{\n"

# Map each element to ID in table
cpp_string += "    static std::map<std::string, int> const ion_map_ids = {"
for count, name in enumerate(ion_names):
    cpp_string += '\n        {"' + name + '", ' + str(count) + "},"
cpp_string = cpp_string[:-1]
cpp_string += " };\n\n"

# Atomic number of each species
cpp_string += "    constexpr int nelements = " + str(len(ion_names)) + ";\n\n"
cpp_string += "    constexpr int ion_atomic_numbers[nelements] = {\n        "
for count, atom_num in enumerate(ion_atom_numbers):
    if count % 10 == 0 and count > 0:
        cpp_string = cpp_string[:-2] + ",\n        "
    cpp_string += str(atom_num) + ", "
cpp_string = cpp_string[:-2]
cpp_string += "};\n\n"

# Offset of each element in table of ionization energies
cpp_string += "    constexpr int ion_energy_offsets[nelements] = {\n        "
for count, offset in enumerate(ion_offsets):
    if count % 10 == 0 and count > 0:
        cpp_string = cpp_string[:-2] + ",\n        "
    cpp_string += str(offset) + ", "
cpp_string = cpp_string[:-2]
cpp_string += "};\n\n"

# Table of ionization energies
cpp_string += (
    "    constexpr int energies_tab_length = " + str(len(list_of_tuples)) + ";\n\n"
)
cpp_string += (
    "    constexpr amrex::Real table_ionization_energies[energies_tab_length]{"
)
for element in ion_names:
    cpp_string += "\n        // " + element + "\n        "
    regex_command = (
        "\n\s+(\d+)\s+\|\s+%s\s+\w+\s+\|\s+\+*(\d+)\s+\|\s+\(*\[*(\d+\.*\d*)" % element
    )
    list_of_tuples = re.findall(regex_command, text_data)
    for count, energy in enumerate([x[2] for x in list_of_tuples]):
        if count % 3 == 0 and count > 0:
            cpp_string = cpp_string[:-2] + ",\n        "
        cpp_string += "amrex::Real(" + energy + "), "
    cpp_string = cpp_string[:-1]
cpp_string = cpp_string[:-1]
cpp_string += "\n    };\n\n"

cpp_string += "}\n\n"

# Write the string to file
cpp_string += "#endif // #ifndef WARPX_UTILS_PHYSICS_IONIZATION_TABLE_H_\n"
f = open("IonizationEnergiesTable.H", "w")
f.write(cpp_string)
f.close()
